from .config import CfgNode as CN

# -----------------------------------------------------------------------------
# Convention about Training / Test specific parameters
# -----------------------------------------------------------------------------
# Whenever an argument can be either used for training or for testing, the
# corresponding name will be post-fixed by a _TRAIN for a training parameter,
# or _TEST for a test-specific parameter.
# For example, the number of images during training will be
# IMAGES_PER_BATCH_TRAIN, while the number of images for testing will be
# IMAGES_PER_BATCH_TEST

# -----------------------------------------------------------------------------
# Config definition
# -----------------------------------------------------------------------------

_C = CN()

# -----------------------------------------------------------------------------
# MODEL
# -----------------------------------------------------------------------------
_C.MODEL = CN()
_C.MODEL.DEVICE = "cuda"
_C.MODEL.META_ARCHITECTURE = 'Baseline'

# ---------------------------------------------------------------------------- #
# Backbone options
# ---------------------------------------------------------------------------- #
_C.MODEL.BACKBONE = CN()
_C.MODEL.BACKBONE.PRETRAINED_PATH = "../models/bert-base-chinese"


# ---------------------------------------------------------------------------- #
# HEADS options
# ---------------------------------------------------------------------------- #
_C.MODEL.HEADS = CN()
_C.MODEL.HEADS.NAME = "TwoLinear"


# ---------------------------------------------------------------------------- #
# LOSSES options
# ---------------------------------------------------------------------------- #
_C.MODEL.LOSSES = CN()
_C.MODEL.LOSSES.NAME = ("CrossEntropyLoss",)

# ---------------------------------------------------------------------------- #
# TOKENIZER options
# ---------------------------------------------------------------------------- #
_C.MODEL.TOKENIZER = CN()
_C.MODEL.TOKENIZER.PRETRAINED_PATH = "../models/bert-base-chinese"
_C.MODEL.TOKENIZER. MAX_LEN = 512

# -----------------------------------------------------------------------------
# DataLoader
# -----------------------------------------------------------------------------
_C.DATALOADER = CN()
# P/K Sampler for data loading
_C.DATALOADER.TRAIN_PATH = '../data/train.csv'
# Naive sampler which don't consider balanced identity sampling
_C.DATALOADER.VAL_PATH = '../data/dev.csv'
# Number of instance for each person
_C.DATALOADER.TRAIN_BATCH = 4

# ---------------------------------------------------------------------------- #
# Solver
# ---------------------------------------------------------------------------- #
_C.SOLVER = CN()

_C.SOLVER.MAX_EPOCHES = 60

_C.SOLVER.BASE_LR = 1e-05

# AUTOMATIC MIXED PRECISION
_C.SOLVER.AMP_ENABLED = False



# ---------------------------------------------------------------------------- #
# Misc options
# ---------------------------------------------------------------------------- #
_C.OUTPUT_DIR = '../logs/test'

# Benchmark different cudnn algorithms.
# If input images have very different sizes, this option will have large overhead
# for about 10k iterations. It usually hurts total time, but can benefit for certain models.
# If input images have the same or similar sizes, benchmark is often helpful.
_C.CUDNN_BENCHMARK = False

